import argparse
import os
import sys
import time
import warnings
from piq import ssim
import math
import torch
import torch.nn.functional as F
from PIL import Image
from pytorch_msssim import ms_ssim
from torchvision import transforms
from torchvision.transforms import ToPILImage
# from .models import TCM_residual_wave_two_entropy_modified_y_downsample_8
from models import TCM_residual_wave_two_entropy_modified_y_downsample_8
import pandas as pd

warnings.filterwarnings("ignore")

print(torch.cuda.is_available())

def save_experiments_to_excel(data, columns, save_path="result.xlsx"):
    df = pd.DataFrame(data=data, columns=columns)
    df.to_excel(save_path, index=False)


def toPilImage(x):
    return ToPILImage()(x.clamp_(0, 1).squeeze())

def compute_psnr(a, b):
    mse = torch.mean((a - b)**2).item()
    return -10 * math.log10(mse)

def compute_msssim(a, b):
    return -10 * math.log10(1-ms_ssim(a, b, data_range=1.).item())

def compute_ssim(a,b):
    return ssim(a, b, data_range=1.).item()

def compute_bpp(out_net):
    size = out_net['x_hat'].size()
    num_pixels = size[0] * size[2] * size[3]
    return sum(torch.log(likelihoods).sum() / (-math.log(2) * num_pixels)
              for likelihoods in out_net['likelihoods'].values()).item()

def pad(x, p):
    h, w = x.size(2), x.size(3)
    new_h = (h + p - 1) // p * p
    new_w = (w + p - 1) // p * p
    padding_left = (new_w - w) // 2
    padding_right = new_w - w - padding_left
    padding_top = (new_h - h) // 2
    padding_bottom = new_h - h - padding_top
    x_padded = F.pad(
        x,
        (padding_left, padding_right, padding_top, padding_bottom),
        mode="constant",
        value=0,
    )
    return x_padded, (padding_left, padding_right, padding_top, padding_bottom)

def crop(x, padding):
    return F.pad(
        x,
        (-padding[0], -padding[1], -padding[2], -padding[3]),
    )

def parse_args(argv):
    parser = argparse.ArgumentParser(description="Example testing script.")
    parser.add_argument("--cuda", action="store_true", help="Use cuda")
    parser.add_argument(
        "--clip_max_norm",
        default=1.0,
        type=float,
        help="gradient clipping max norm (default: %(default)s",
    )
    parser.add_argument("--checkpoint", type=str, help="Path to a checkpoint")
    parser.add_argument("--N", type=int,default=128, help="The number filters")
    parser.add_argument("--data", type=str, help="Path to dataset")
    parser.add_argument(
        "--real", action="store_true", default=True
    )
    parser.set_defaults(real=False)
    args = parser.parse_args(argv)
    return args


def main(argv):
    data = []
    columns = ["image", "PSNR", "MS-SSIM", "Bpp"]
    args = parse_args(argv)
    p = 128
    path = args.data
    img_list = []
    for file in os.listdir(path):
        if file[-3:] in ["jpg", "png", "peg"]:
            img_list.append(file)
    if args.cuda:
        device = 'cuda:0'
    else:
        device = 'cpu'
    net = TCM_residual_wave_two_entropy_modified_y_downsample_8(config=[2,2,2,2,2,2], head_dim=[8, 16, 32, 32, 16, 8], drop_path_rate=0.0, N=args.N, M=320)
    net = net.to(device)
    net.eval()
    count = 0
    PSNR = 0
    Bit_rate = 0
    MS_SSIM = 0
    total_time = 0
    total_enc =0
    total_dec =0
    dictory = {}
    if args.checkpoint:  # load from previous checkpoint
        print("Loading", args.checkpoint)
        checkpoint = torch.load(args.checkpoint, map_location=device)
        for k, v in checkpoint["state_dict"].items():
            dictory[k.replace("module.", "")] = v
        net.load_state_dict(dictory)
    if args.real:
        net.update()
        for img_name in img_list:
            print(img_name)
            img_path = os.path.join(path, img_name)
            img = transforms.ToTensor()(Image.open(img_path).convert('RGB')).to(device)
            x = img.unsqueeze(0)
            x_padded, padding = pad(x, p)
            count += 1
            if count > 3: break
            with torch.no_grad():
                if args.cuda:
                    torch.cuda.synchronize()
                s = time.time()
                out_enc = net.compress(x_padded)
                enc_time = time.time()-s
                dec_start = time.time()
                out_dec = net.decompress(out_enc["strings"], out_enc["shape"])
                if args.cuda:
                    torch.cuda.synchronize()
                e = time.time()
                dec_time = e-dec_start
                total_time += (e - s)
                total_enc +=enc_time
                total_dec +=dec_time
                out_dec["x_hat"] = crop(out_dec["x_hat"], padding)

                num_pixels = x.size(0) * x.size(2) * x.size(3)

                # NOTE: 
                img = toPilImage(out_dec["x_hat"])
                folder = r"F:\PCL月报\月报\2025年\25-06月\实验记录\WeConvene\CLIC\0.05"
                os.makedirs(folder, exist_ok=True)
                img_save_path = os.path.join(folder, img_name.replace(".png", "_dec.png"))
                img.save(img_save_path)

                print(f'Bitrate: {(sum(len(s[0]) for s in out_enc["strings"]) * 8.0 / num_pixels):.3f}bpp')
                print(f'SSIM: {compute_ssim(x, out_dec["x_hat"]):.2f}')
                print(f'PSNR: {compute_psnr(x, out_dec["x_hat"]):.2f}dB')
                Bit_rate += sum(len(s[0]) for s in out_enc["strings"]) * 8.0 / num_pixels
                PSNR += compute_psnr(x, out_dec["x_hat"])
                MS_SSIM += compute_msssim(x, out_dec["x_hat"])

                # NOTE: 
                data.append(
                    [
                    img_name,
                    compute_psnr(x, out_dec["x_hat"]), 
                    compute_ssim(x, out_dec["x_hat"]),
                    (sum(len(s[0]) for s in out_enc["strings"]) * 8.0 / num_pixels)
                    ]
                )

    else:
        for img_name in img_list:
            print(img_name)
            img_path = os.path.join(path, img_name)
            img = Image.open(img_path).convert('RGB')
            x = transforms.ToTensor()(img).unsqueeze(0).to(device)
            x_padded, padding = pad(x, p)
            count += 1
            with torch.no_grad():
                if args.cuda:
                    torch.cuda.synchronize()
                s = time.time()
                out_net = net.forward(x_padded)
                if args.cuda:
                    torch.cuda.synchronize()
                e = time.time()
                total_time += (e - s)
                out_net['x_hat'].clamp_(0, 1)
                out_net["x_hat"] = crop(out_net["x_hat"], padding)
                print(f'PSNR: {compute_psnr(x, out_net["x_hat"]):.2f}dB')
                print(f'SSIM: {compute_ssim(x, out_net["x_hat"]):.2f}')
                print(f'Bit-rate: {compute_bpp(out_net):.3f}bpp')
                PSNR += compute_psnr(x, out_net["x_hat"])
                MS_SSIM += compute_ssim(x, out_net["x_hat"])
                Bit_rate += compute_bpp(out_net)
    PSNR = PSNR / count
    MS_SSIM = MS_SSIM / count
    Bit_rate = Bit_rate / count
    total_time = total_time / count
    ave_enc  = total_enc / count
    ave_dec  = total_dec /count
    print(f'average_PSNR: {PSNR:.2f}dB')
    print(f'average_SSIM: {MS_SSIM:.4f}')
    print(f'average_Bit-rate: {Bit_rate:.3f} bpp')
    print(f'average_time: {total_time:.3f} ms')
    print(f'encode average_time: {ave_enc:.3f} ms')
    print(f'decode average_time: {ave_dec:.3f} ms')

    # NOTE: 
    save_path = r"F:\PCL月报\月报\2025年\25-05月\实验记录\WeConvene\Tecnick_RGB_1200x1200-图像解码\MSE_0.05_CLIC.xlsx"
    save_experiments_to_excel(data=data, columns=columns, save_path=save_path)
    

if __name__ == "__main__":
    print(torch.cuda.is_available())
    main(sys.argv[1:])

"""
python eval_wave_residual_two_entropy_modified_y_downsample_8.py --checkpoint "D:/BaiduNetdiskDownload/ECCV2024/0.05/0.05175_checkpoint.pth.tar" --data D:/video-communication-dataset/图像编解码测试/CLIC/professional_valid/ --cuda --real
"""